/** \file    refine/smooth.cpp
  * \brief   Construst smooth grey mask from existing bit mask for the object.
  * \date    2004
  * \author  Fedor Chelnokov
  * \version $Header: /home/cscom1r4/cvsroot/cs/refine/smooth.cpp,v 1.8 2005/06/26 11:20:46 cscom1r4 Exp $
  */

#include <limits>
#include <iostream>
#include "refine/smooth.h"
#include "imglib/imageclass.h"
#include "imglib/bitmap.h"
#include "imglib/stat_calculator.h"
#include "imglib/dichromatic_raster.h"
#include "imglib/bmpmorph.h"
#include "refine/topology_limit.h"
#include "xtd/minmax"

using std::clog;
using std::endl;

namespace refine {

namespace smooth {

const params default_params = {5, false, info_mean, false, false, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;

template <class R, class M>
class smoother {
public:
	typedef typename R::elem_type color_type;
public:
	const params & _params;
public:
	smoother(const params & p, const R & s, const bit_mask::mask & m) 
		: _params(p), _src(s), _src_mask(&m) 
		{ }
	void run(M & b);

	const R & src() const
		{ return _src; }
	const color_type & src(size_point p) const
		{ return _src[p]; }
	const color_type & src(int_point p) const
		{ return _src[p]; }
	const color_type & src(unsigned x, unsigned y) const
		{ return _src(x,y); }
	unsigned width() const
		{ return _src.width(); }
	unsigned height() const
		{ return _src.height(); }
	size_extent dim() const
		{ return _src.dim(); }

	const bit_mask::mask & src_mask() const
		{ return *_src_mask; }
	bool src_mask(size_point p) const
		{ return (*_src_mask)[p]; }
	bool src_mask(int_point p) const
		{ return (*_src_mask)[p]; }
	bool src_mask(unsigned x, unsigned y) const
		{ return (*_src_mask)(x,y); }

	const M & smoothed() const
		{ return *_smoothed; }
	grey_mask_value smoothed(size_point p) const
		{ return (*_smoothed)[p]; }
	grey_mask_value smoothed(int_point p) const
		{ return (*_smoothed)[p]; }
	grey_mask_value smoothed(unsigned x, unsigned y) const
		{ return (*_smoothed)(x,y); }
private:
	///\return true if a change was made
	bool test(size_point p);
private:
	const R & _src;
	const bit_mask::mask * _src_mask;
	M * _smoothed;
	int_rect _frame;
};

template <class R, class M>
static void perform_impl(const R & src, const bit_mask::mask & m, M & s, const params & p)
{
	if (!p.has_initial_apprx)
		s = img::make_dichromatic_raster(m, grey_mask_min_value, grey_mask_max_value);
	smoother<R, M>(p, src, m).run(s);
}

void perform(const byte_image & src, const bit_mask::mask & m, grey_mask::mask & s, const params & p)
	{ perform_impl(src, m, s, p); }

void perform(const argb_byte_image & src, const bit_mask::mask & m, grey_mask::mask & s, const params & p)
	{ perform_impl(src, m, s, p); }

void perform(byte_image & img, const bit_mask::mask & m, const params & p)
	{ perform_impl(img, m, img::projection(img, &byte_color::a), p); }

void perform(argb_byte_image & img, const bit_mask::mask & m, const params & p)
	{ perform_impl(img, m, img::projection(img, &argb_byte_color::a), p); }

// smoother implementation //

template <class R, class M>
void smoother<R, M>::run(M & b)
{
	_smoothed = & b;
	_frame = int_rect(0, 0, width(), height());

	assert(src_mask().dim() == dim());
	assert(_smoothed->dim() == dim());

	img::bitmap fixed = *_src_mask; // alpha for fixed points will not change
	img::erode8(fixed, _params.window_half_size);
	{
		img::bitmap tmp = *_src_mask;
		tmp.inverse();
		img::erode8(tmp, _params.window_half_size);
		img::combine(fixed, tmp, img::bit_or());
	}

	img::bitmap curr(dim()), next(dim());
	curr.fill(true);

	if (!_params.silent)
		clog << "mask smoothing initiated..." << endl;

	for (int pass = 0; pass < 150; ++pass, curr.swap(next)) {
		unsigned tests = 0, changes = 0;

		next.fill(false);
		img::bitmap::iterator i;
		for (i = curr.begin(); i != curr.end(); ++i) {
			if (fixed[i] || !*i)
				continue;
			*i = false;
			size_point p = i;
			if (p.x == 0 || p.y == 0 || p.x == width()-1 || p.y == height()-1)
				continue;
			++tests;
			if (test(p)) {
				++changes;
				next[p - size_point(-1,1)] = true;
				next[p - size_point(0,1)] = true;
				next[p - size_point(1,1)] = true;
				next[p - size_point(1,0)] = true;
				curr[p + size_point(-1,1)] = true;
				curr[p + size_point(0,1)] = true;
				curr[p + size_point(1,1)] = true;
				curr[p + size_point(1,0)] = true;
			}
		}

		assert(curr.set_bit_count() == 0);

		if (!_params.silent)
			clog << "pass " << pass <<
					": tests=" << tests << 
					", changes=" << changes << 
					endl;

		if (!_params.monotone || changes == 0)
			break;
	}

	if (!_params.silent)
		clog << endl;
}

struct color_distance {
	color_distance() { }
	color_distance(byte_color const & c, byte_color const & e) : color(c), 
		distance((img::BGR<int>(c)-img::BGR<int>(e)).sqr<unsigned>()) { }

	byte_color color;
	unsigned distance;
};

bool operator < (color_distance const & lhs, color_distance const & rhs)
	{ return lhs.distance < rhs.distance; }

template <class R, class M>
bool smoother<R, M>::test(size_point p)
{
	grey_mask_value min_nei, max_nei;
	if (_params.monotone) {
		refine::grey_mask::topology_limit(smoothed(), p, min_nei, max_nei);
		if (min_nei > max_nei)
			return false;
	}
	else {
		min_nei = grey_mask_min_value;
		max_nei = grey_mask_max_value;
	}

	grey_mask_value v0 = smoothed(p);

	if (min_nei == max_nei) {
		if (v0 != min_nei) {
			(*_smoothed)[p] = min_nei;
			return true;
		}
		return false;
	}

	img::stat_calculator calc[2];

	if (_params.cluster) {
		std::vector<color_distance> v[2];

		color_type c = src(p);
		int_point d;
		for (d.y = -_params.window_half_size; d.y <= _params.window_half_size; ++d.y) {
			for (d.x = -_params.window_half_size; d.x <= _params.window_half_size; ++d.x) {
				int_point z = int_point(p) + d;
				if (!_frame.inside(z))
					continue;
				int i = refine::ii(src_mask(z));
				v[i].push_back(color_distance(src(z), c));
			}
		}
		if (v[0].size() < 7) {
			calc[1].add(v[1][0].color);
		}
		else if (v[1].size() < 7) {
			calc[0].add(v[0][0].color);
		}
		else {
			size_t n = std::min(v[0].size(), v[1].size()) / 3;
			for (int i = 0; i < 2; ++i) {
				std::partial_sort(v[i].begin(), v[i].begin() + n, v[i].end());
				for (int j = 0; j < n; ++j)
					calc[i].add(v[i][j].color);
			}
		}
	}
	else {
		int_point d;
		for (d.y = -_params.window_half_size; d.y <= _params.window_half_size; ++d.y) {
			for (d.x = -_params.window_half_size; d.x <= _params.window_half_size; ++d.x) {
				int_point z = int_point(p) + d;
				if (!_frame.inside(z))
					continue;
				int i = refine::ii(src_mask(z));
				calc[i].add(src(z));
			}
		}
	}

	grey_mask_value vx;
	if (calc[0].count() == 0)
		vx = grey_mask_max_value;
	else if (calc[1].count() == 0)
		vx = grey_mask_min_value;
	else {
		img::BGR<float> m[2];
		m[0] = calc[0].mean();
		m[1] = calc[1].mean();

		//find position of current color relative to the oriented cut(m[0]->m[1]):
		//coordinate of color's projection on the cut, where a = 0 for m[0], a = 1 for m[1]
		/*float_color g = m[1] - m[0];
		float_color c = float_color(src(p)) - m[0];
		float a = (g.r*c.r + g.g*c.g + g.b*c.b) / g.sqr();

		if (a <= 0)
			vx = grey_mask_min_value;
		else if (a >= 1)
			vx = grey_mask_max_value;
		else
			vx = linear(grey_mask_min_value, grey_mask_max_value, a);*/

		img::BGR<float> c(src(p));
		float mm = (m[0] - m[1]).sqr<float>();
		float rr[2];
		rr[0] = (c - m[0]).sqr<float>();
		rr[1] = (c - m[1]).sqr<float>();

		if (_params.consider == info_mean_and_dispersion) {
			float ss[2]; //dispersions
			ss[0] = calc[0].dispersion();
			ss[1] = calc[1].dispersion();
			bool low[2];
			low[0] = ss[0] < std::numeric_limits<float>::epsilon();
			low[1] = ss[1] < std::numeric_limits<float>::epsilon();
			if (low[0] && low[1]) 
				; //nothing
			else if (low[0]) {
				rr[0] = 1;
				rr[1] = 0;
			}
			else if (low[1]) {
				rr[0] = 0;
				rr[1] = 1;
			}
			else {
				rr[0] /= ss[0];
				rr[1] /= ss[1];
				mm /= sqrt(ss[0] * ss[1]);
			}
		}

		if (rr[0] == rr[1])
			vx = grey_mask_avg_value;
		else if (rr[0] < std::numeric_limits<float>::epsilon() || rr[1] > mm + rr[0])
			vx = grey_mask_min_value;
		else if (rr[1] < std::numeric_limits<float>::epsilon() || rr[0] > mm + rr[1])
			vx = grey_mask_max_value;
		else {
			float r[2];
			r[0] = 1.0 / sqrt(rr[0]);
			r[1] = 1.0 / sqrt(rr[1]);
			vx = linear(grey_mask_min_value, grey_mask_max_value, r[1] / (r[0] + r[1]));
		}
	}

	xtd::max(vx, min_nei);
	xtd::min(vx, max_nei);

	if (v0 == vx)
		return false;

	(*_smoothed)[p] = vx;
	return true;
}

} //namespace smooth

} //namespace refine
